import os
import random
import shutil

# 设置原始文件夹路径和目标文件夹路径
original_folder = r"F:\new_deep_learning\data\flowers"
train_folder = "trains"
test_folder = "tests"

# 创建目标文件夹
os.makedirs(train_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)

# 获取原始文件夹中的子文件夹列表
subfolders = [f.path for f in os.scandir(original_folder) if f.is_dir()]

# 遍历每个子文件夹
for subfolder in subfolders:
	# 获取子文件夹名称
	folder_name = os.path.basename(subfolder)
	
	# 创建目标文件夹的子文件夹
	train_subfolder = os.path.join(train_folder, folder_name)
	test_subfolder = os.path.join(test_folder, folder_name)
	os.makedirs(train_subfolder, exist_ok=True)
	os.makedirs(test_subfolder, exist_ok=True)
	
	# 获取子文件夹中的图片文件列表
	images = [f.path for f in os.scandir(subfolder) if f.is_file()]
	
	# 计算训练集数量和测试集数量
	total_images = len(images)
	train_count = int(total_images * 0.8)
	test_count = total_images - train_count
	
	# 打乱图片顺序
	random.shuffle(images)
	
	# 将图片复制到训练集文件夹
	for i in range(train_count):
		image = images[i]
		shutil.copy(image, train_subfolder)
	
	# 将图片复制到测试集文件夹
	for i in range(train_count, total_images):
		image = images[i]
		shutil.copy(image, test_subfolder)

print("数据集划分完成！")
